// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/gtest_utils.h>
#include <pollux/common/testutil/optional_empty.h>
#include <pollux/testing/dwio/batch_maker.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;

class MapFilterTest : public functions::test::FunctionBaseTest {
 protected:
  template <typename K, typename V>
  void checkMapFilter(
      BaseVector* inputMap,
      const BaseVector& result,
      std::function<bool(K*, V*, vector_size_t, vector_size_t)> test) {
    auto resultMap = result.wrapped_vector()->as<MapVector>();
    auto resultKeys = resultMap->mapKeys()->as<K>();
    auto resultValues = resultMap->mapValues()->as<V>();
    auto map = inputMap->wrapped_vector()->as<MapVector>();
    auto keys = map->mapKeys()->as<K>();
    auto values = map->mapValues()->as<V>();
    for (auto i = 0; i < inputMap->size(); ++i) {
      bool expectedNull = inputMap->is_null_at(i);
      EXPECT_EQ(expectedNull, result.is_null_at(i));
      if (expectedNull) {
        continue;
      }
      auto mapIndex = inputMap->wrapped_index(i);
      auto offset = map->offsetAt(mapIndex);
      auto size = map->sizeAt(mapIndex);

      auto resultIndex = result.wrapped_index(i);
      auto resultOffset = resultMap->offsetAt(resultIndex);
      auto resultSize = resultMap->sizeAt(resultIndex);
      int32_t count = 0;
      for (auto j = offset; j < offset + size; ++j) {
        if (test(keys, values, j, i)) {
          auto resultElementIndex = resultOffset + count;
          ++count;
          EXPECT_LE(count, resultSize)
              << "at " << i << ": " << resultMap->toString(resultIndex);
          EXPECT_TRUE(keys->equal_value_at(resultKeys, j, resultElementIndex))
              << "at (" << i << ", " << j << "): " << keys->toString(j)
              << " vs. " << resultKeys->toString(resultElementIndex);
          EXPECT_TRUE(values->equal_value_at(resultValues, j, resultElementIndex))
              << "at (" << i << ", " << j << "): " << values->toString(j)
              << " vs. " << resultValues->toString(resultElementIndex);
        }
      }
      EXPECT_EQ(resultSize, count)
          << "at " << i << ": " << resultMap->toString(resultIndex);
    }
  }
};

TEST_F(MapFilterTest, filter) {
  auto rowType =
      ROW({"long_val", "map_val"}, {BIGINT(), MAP(BIGINT(), INTEGER())});
  auto data = std::static_pointer_cast<RowVector>(
      BatchMaker::createBatch(rowType, 1'000, *execCtx_.pool()));

  auto result = evaluate("map_filter(map_val, (k, v) -> (k > long_val))", data);
  auto* cutoff = data->childAt(0)->as<SimpleVector<int64_t>>();
  checkMapFilter<SimpleVector<int64_t>, SimpleVector<int32_t>>(
      data->childAt(1).get(),
      *result,
      [&](SimpleVector<int64_t>* keys,
          SimpleVector<int32_t>* values,
          vector_size_t elementRow,
          vector_size_t row) {
        return cutoff->is_null_at(row) || keys->is_null_at(elementRow)
            ? false
            : keys->value_at(elementRow) > cutoff->value_at(row);
      });
}

TEST_F(MapFilterTest, empty) {
  auto rowType =
      ROW({"long_val", "map_val"}, {BIGINT(), MAP(BIGINT(), INTEGER())});
  auto input = std::static_pointer_cast<RowVector>(
      BatchMaker::createBatch(rowType, 1'000, *execCtx_.pool()));

  auto result =
      evaluate<MapVector>("map_filter(map_val, (k, v) -> (k = 11111))", input);

  EXPECT_EQ(result->size(), input->size());
  auto inputMap = input->childAt(1);
  for (auto i = 0; i < input->size(); ++i) {
    bool isNull = inputMap->is_null_at(i);
    EXPECT_EQ(isNull, result->is_null_at(i)) << "at " << i;
    if (!isNull) {
      EXPECT_EQ(result->sizeAt(i), 0)
          << "at " << i << ": " << result->toString(i);
    }
  }
}

TEST_F(MapFilterTest, dictionaryWithUniqueValues) {
  auto rowType =
      ROW({"long_val", "map_val"}, {BIGINT(), MAP(BIGINT(), INTEGER())});
  auto data = std::static_pointer_cast<RowVector>(
      BatchMaker::createBatch(rowType, 10, *execCtx_.pool()));

  // Wrap the input in a dictionary.
  BufferPtr indices = make_indices_in_reverse(data->size());
  data->childAt(1) = wrap_in_dictionary(indices, data->size(), data->childAt(1));
  auto result = evaluate(
      "map_filter(map_filter(map_val, (k, v) -> (k > long_val)), (k, v) -> (v > 0))",
      data);
  auto* cutoff = data->childAt(0)->as<SimpleVector<int64_t>>();
  auto test = [&](SimpleVector<int64_t>* keys,
                  SimpleVector<int32_t>* values,
                  vector_size_t elementRow,
                  vector_size_t row) {
    if (cutoff->is_null_at(row) || keys->is_null_at(elementRow) ||
        values->is_null_at(elementRow)) {
      return false;
    }
    auto cutoffValue = cutoff->value_at(row);
    return keys->value_at(elementRow) > cutoffValue &&
        values->value_at(elementRow) > 0;
  };
  checkMapFilter<SimpleVector<int64_t>, SimpleVector<int32_t>>(
      data->childAt(1).get(), *result, test);

  // Wrap both inputs in the same dictionary.
  data->childAt(0) = wrap_in_dictionary(indices, data->size(), data->childAt(0));
  result = evaluate(
      "map_filter(map_filter(map_val, (k, v) -> (k > long_val)), (k, v) -> (v > 0))",
      data);
  cutoff = data->childAt(0)->as<SimpleVector<int64_t>>();
  checkMapFilter<SimpleVector<int64_t>, SimpleVector<int32_t>>(
      data->childAt(1).get(), *result, test);
}

TEST_F(MapFilterTest, conditional) {
  auto rowType =
      ROW({"long_val", "map_val"}, {BIGINT(), MAP(BIGINT(), INTEGER())});
  auto data = std::static_pointer_cast<RowVector>(
      BatchMaker::createBatch(rowType, 1'000, *execCtx_.pool()));

  auto result = evaluate(
      "map_filter(map_val, "
      "  if (long_val < 0, (k, v) -> (v < long_val), (k, v) -> (k > long_val)))",
      data);

  auto* cutoff = data->childAt(0)->as<SimpleVector<int64_t>>();
  auto test = [&](SimpleVector<int64_t>* keys,
                  SimpleVector<int32_t>* values,
                  vector_size_t elementRow,
                  vector_size_t row) {
    if (cutoff->is_null_at(row)) {
      return false;
    }
    auto cutoffValue = cutoff->value_at(row);
    if (cutoffValue < 0) {
      return keys->is_null_at(elementRow)
          ? false
          : keys->value_at(elementRow) < cutoffValue;
    } else {
      return values->is_null_at(elementRow)
          ? false
          : values->value_at(elementRow) > cutoffValue;
    }
  };

  checkMapFilter<SimpleVector<int64_t>, SimpleVector<int32_t>>(
      data->childAt(1).get(), *result, test);
}

TEST_F(MapFilterTest, dictionaryWithDuplicates) {
  vector_size_t size = 1'000;

  // make a map vector where each row repeats a few times
  auto sizeAt = [](vector_size_t row) { return row % 5; };
  auto baseMap = make_map_vector<int32_t, int64_t>(
      size / 2,
      sizeAt,
      [](vector_size_t row) { return row % 7; },
      [](vector_size_t row) { return row % 11; },
      null_every(11));

  BufferPtr indices = make_indices(size, [](auto row) { return row / 2; });

  auto map = wrap_in_dictionary(indices, size, baseMap);

  // make a capture with unique values
  auto capture =
      make_flat_vector<int32_t>(size, [](vector_size_t row) { return row; });

  auto input = make_row_vector({capture, map});

  auto result =
      evaluate("map_filter(c1, (k, v) -> ((k + v + c0) % 7 < 3))", input);

  auto flatMap = flatten(map);
  input = make_row_vector({capture, flatMap});
  auto expectedResult =
      evaluate("map_filter(c1, (k, v) -> ((k + v + c0) % 7 < 3))", input);

  assertEqualVectors(expectedResult, result);
}

TEST_F(MapFilterTest, lambdaSelectivityVector) {
  auto data = make_row_vector({
      wrap_in_dictionary(
          make_indices({0}),
          1,
          make_flat_vector<int64_t>(std::vector<int64_t>{10})),
  });

  // Our expression. Use large numbers to trigger asan if things go wrong.
  auto exprSet = compileExpression(
      "map_filter("
      "MAP(ARRAY[233439836560246536, 398885052601874414, 213334509704047604],"
      "ARRAY[c0, c0, c0]),"
      "(k, v) -> (v IS NOT NULL))",
      as_row_type(data->type()));

  // Ensure that our context would have 'final selection' false.
  exec::EvalCtx context(&execCtx_, exprSet.get(), data.get());
  const SelectivityVector allRows(data->size());
  exec::ScopedFinalSelectionSetter scopedFinalSelectionSetter(
      context, &allRows);

  // Evaluate. Result would be overwritten.
  std::vector<VectorPtr> result = {
      make_flat_vector<int64_t>(std::vector<int64_t>{1})};
  exprSet->eval(allRows, context, result);

  auto expectedKeys = make_flat_vector<int64_t>(
      {233439836560246536, 398885052601874414, 213334509704047604});
  auto expectedValues = make_flat_vector<int64_t>({10, 10, 10});
  auto expected = make_map_vector({0}, expectedKeys, expectedValues);
  assertEqualVectors(expected, result[0]);
}

TEST_F(MapFilterTest, try) {
  auto data = make_row_vector({
      make_map_vector<int64_t, int64_t>({
          {{1, 2}, {2, 3}},
          {{3, 4}, {0, 1}, {5, 6}},
          {{6, 5}, {7, 8}},
          {{8, 7}},
      }),
  });

  POLLUX_ASSERT_THROW(
      evaluate("map_filter(c0, (k, v) -> (v / k > 0))", data),
      "division by zero");

  auto result = evaluate("try(map_filter(c0, (k, v) -> (v / k > 0)))", data);
  auto expected = make_nullable_map_vector<int64_t, int64_t>(
      {{{{1, 2}, {2, 3}}},
       std::nullopt,
       {{{7, 8}}},
       common::testutil::optionalEmpty});
  assertEqualVectors(expected, result);
}

TEST_F(MapFilterTest, unknown) {
  auto data = make_row_vector({make_all_null_map_vector(10, UNKNOWN(), BIGINT())});
  auto result = evaluate("map_filter(c0, (k, v) -> (v > 5))", data);
  assertEqualVectors(data->childAt(0), result);
}

TEST_F(MapFilterTest, selectiveFilter) {
  // Verify that a selective filter will ensure the underlying elements
  // vector is flattened before generating the result which is otherwise wrapped
  // in a dictionary with the filter results. This ensures large element
  // vectors are not passed along.
  auto data = make_row_vector({
      make_map_vector<int64_t, int64_t>(
          {{{1, 3},
            {2, 3},
            {3, 3},
            {4, 3},
            {5, 3},
            {6, 3},
            {7, 3},
            {8, 3},
            {9, 3},
            {10, 3},
            {11, 3},
            {12, 3},
            {13, 3},
            {14, 3},
            {15, 3},
            {16, 3}}}),
  });

  auto result = evaluate("map_filter(c0, (k, v) -> (k = 1))", data);
  auto base = result->as<MapVector>()->mapKeys();
  EXPECT_EQ(base->encoding(), VectorEncoding::Simple::FLAT);
  base = result->as<MapVector>()->mapValues();
  EXPECT_EQ(base->encoding(), VectorEncoding::Simple::FLAT);

  result = evaluate("map_filter(c0, (k, v) -> (k < 6))", data);
  base = result->as<MapVector>()->mapKeys();
  EXPECT_EQ(base->encoding(), VectorEncoding::Simple::DICTIONARY);
  base = result->as<MapVector>()->mapValues();
  EXPECT_EQ(base->encoding(), VectorEncoding::Simple::DICTIONARY);
}
